﻿using System;
using System.Buffers;
using System.Buffers.Binary;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace WholeSlideImageLibrary.Formats.OptimizedPyramidTiledImage
{
    public class OptimizedPyramidTiledImageWriter
    {
        private Stream _stream;
        private bool _leaveOpen;
        private long _position;

        public long BytesWritten => _position;

        public OptimizedPyramidTiledImageWriter(Stream stream, bool leaveOpen)
        {
            _stream = stream ?? throw new ArgumentNullException(nameof(stream));
            _leaveOpen = leaveOpen;
            _position = 0;
        }

        public void Seek(long position)
        {
            _stream.Seek(position, SeekOrigin.Begin);
        }

        public void SeekToEnd()
        {
            _stream.Seek(_position, SeekOrigin.Begin);
        }

        public readonly static int FileHeaderLength = 32;

        public async Task<int> ReadToAsync(Stream stream, int length, CancellationToken cancellationToken = default)
        {
            if (stream is null)
            {
                throw new ArgumentNullException(nameof(stream));
            }

            int bytesCoppied = 0;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(Math.Min(length, 81920));
            try
            {
                while (bytesCoppied < length)
                {
                    int bytesRead = await _stream.ReadAsync(buffer, 0, Math.Min(length - bytesCoppied, buffer.Length), cancellationToken).ConfigureAwait(false);
                    if (bytesRead == 0)
                    {
                        break;
                    }
                    await stream.WriteAsync(buffer, 0, bytesRead, cancellationToken).ConfigureAwait(false);
                    bytesCoppied += bytesRead;
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

            return bytesCoppied;
        }

        public async Task WriteFileHeaderAsync(long mainBlockPointer, CancellationToken cancellationToken = default)
        {
            byte[] buffer = ArrayPool<byte>.Shared.Rent(32);
            try
            {
                buffer[0] = (byte)'O';
                buffer[1] = (byte)'P';
                buffer[2] = (byte)'T';
                buffer[3] = (byte)'I';
                buffer[4] = (byte)'0';
                buffer[5] = (byte)'1';
                buffer[6] = 0;
                buffer[7] = 0;
                BinaryPrimitives.WriteInt64LittleEndian(buffer.AsSpan(8), mainBlockPointer);

                Array.Copy(buffer, 0, buffer, 16, 16);

                await _stream.WriteAsync(buffer, 0, 32, cancellationToken).ConfigureAwait(false);
                _position += 32;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
        }

        public async Task<FileContentRegion> WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken = default)
        {
            long position = _position;
            await _stream.WriteAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
            _position = position + count;
            return new FileContentRegion(position, count);
        }

        public async Task<FileContentRegion> WriteAsync(ReadOnlyMemory<byte> memory, CancellationToken cancellationToken = default)
        {
            long position = _position;
            await WriteAsync(_stream, memory, cancellationToken).ConfigureAwait(false);
            _position = position + memory.Length;
            return new FileContentRegion(position, memory.Length);
        }

        public async Task<FileContentRegion> WriteAsync(ReadOnlySequence<byte> sequence, CancellationToken cancellationToken = default)
        {
            long position = _position;
            foreach (ReadOnlyMemory<byte> segment in sequence)
            {
                await WriteAsync(_stream, segment, cancellationToken).ConfigureAwait(false);
                _position += segment.Length;
            }
            return new FileContentRegion(position, (int)sequence.Length);
        }

        private static async Task WriteAsync(Stream stream, ReadOnlyMemory<byte> memory, CancellationToken cancellationToken = default)
        {
            if (MemoryMarshal.TryGetArray(memory, out ArraySegment<byte> arraySegment))
            {
                await stream.WriteAsync(arraySegment.Array!, arraySegment.Offset, arraySegment.Count, cancellationToken).ConfigureAwait(false);
            }
            else
            {
                byte[] buffer = ArrayPool<byte>.Shared.Rent(memory.Length);
                try
                {
                    memory.Span.CopyTo(buffer);
                    await stream.WriteAsync(buffer, 0, memory.Length, cancellationToken).ConfigureAwait(false);
                }
                finally
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                }
            }
        }

        public async Task<FileContentRegion> WriteFromAsync(Stream stream, CancellationToken cancellationToken)
        {
            if (stream is null)
            {
                throw new ArgumentNullException(nameof(stream));
            }

            int bytesCoppied = 0;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(81920);
            try
            {
                while (true)
                {
                    int bytesRead = await stream.ReadAsync(buffer, 0, buffer.Length, cancellationToken).ConfigureAwait(false);
                    if (bytesRead == 0)
                    {
                        break;
                    }
                    await _stream.WriteAsync(buffer, 0, bytesRead, cancellationToken).ConfigureAwait(false);
                    bytesCoppied += bytesRead;
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

            long position = _position;
            _position = position + bytesCoppied;
            return new FileContentRegion(position, bytesCoppied);
        }

        public ValueTask<FileMetadataItem> WriteMetadataItemAsync(string name, string value, CancellationToken cancellationToken = default)
        {
            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }
            if (value is null)
            {
                throw new ArgumentNullException(nameof(value));
            }
            int length = Encoding.UTF8.GetByteCount(value);
            if (length <= 8)
            {
                return new ValueTask<FileMetadataItem>(new FileMetadataItem(name, value));
            }
            return new ValueTask<FileMetadataItem>(SlowWriteStringAsync());

            async Task<FileMetadataItem> SlowWriteStringAsync()
            {
                int len = Encoding.UTF8.GetMaxByteCount(value.Length);
                byte[] buffer = ArrayPool<byte>.Shared.Rent(len);
                try
                {
                    len = Encoding.UTF8.GetBytes(value, 0, value.Length, buffer, 0);
                    await _stream.WriteAsync(buffer, 0, len, cancellationToken).ConfigureAwait(false);
                    long position = _position;
                    _position = position + len;
                    return new FileMetadataItem(name, (ushort)len, position, 0);
                }
                finally
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                }
            }
        }

        public Task<long> WriteMetadataBlockAsync(FileMetadataItem[] items, CancellationToken cancellationToken = default)
            => WriteMetadataBlockAsync(string.Empty, items, 0, cancellationToken);

        public Task<long> WriteMetadataBlockAsync(FileMetadataItem[] items, long nextBlockPointer, CancellationToken cancellationToken = default)
            => WriteMetadataBlockAsync(string.Empty, items, nextBlockPointer, cancellationToken);

        public async Task<long> WriteMetadataBlockAsync(string name, FileMetadataItem[] items, long nextBlockPointer, CancellationToken cancellationToken = default)
        {
            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }

            if (items is null)
            {
                throw new ArgumentNullException(nameof(items));
            }

            int nameLength = Encoding.UTF8.GetByteCount(name);
            if (nameLength > 255)
            {
                throw new ArgumentException("Name too long.");
            }

            int totalLength = 3 + nameLength;
            foreach (FileMetadataItem item in items)
            {
                totalLength += item.GetSize();
            }

            long position = _position;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(4 + totalLength + 8);
            try
            {
                BinaryPrimitives.WriteInt32LittleEndian(buffer.AsSpan(), totalLength);
                buffer[4] = (byte)nameLength;
                EncodingHelper.EncodeUtf8(name, buffer.AsSpan(5, nameLength));
                BinaryPrimitives.WriteInt16LittleEndian(buffer.AsSpan(5 + nameLength), (short)items.Length);
                await _stream.WriteAsync(buffer, 0, 7 + nameLength, cancellationToken).ConfigureAwait(false);
                _position += 7 + nameLength;

                int bytesWritten = 0;
                for (int i = 0; i < items.Length; i++)
                {
                    if (!items[i].TryWriteTo(buffer.AsSpan(bytesWritten), out int written))
                    {
                        throw new InvalidDataException();
                    }
                    bytesWritten += written;
                }
                BinaryPrimitives.WriteInt64LittleEndian(buffer.AsSpan(bytesWritten), nextBlockPointer);
                await _stream.WriteAsync(buffer, 0, bytesWritten + 8, cancellationToken).ConfigureAwait(false);
                _position += bytesWritten + 8;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

            return position;
        }
        public Task<FileAssociatedImageItem> WriteAssociatedImageItemAsync(byte[] buffer, string name, int width, int height, string format, CancellationToken cancellationToken = default)
        {
            if (buffer is null)
            {
                throw new ArgumentNullException(nameof(buffer));
            }

            return WriteAssociatedImageItemAsync(buffer, 0, buffer.Length, name, width, height, format, cancellationToken);
        }

        public async Task<FileAssociatedImageItem> WriteAssociatedImageItemAsync(byte[] buffer, int offset, int count, string name, int width, int height, string format, CancellationToken cancellationToken = default)
        {
            if (format is null || format.Length != 4)
            {
                throw new ArgumentException("Format not supported.", nameof(format));
            }

            await _stream.WriteAsync(buffer, offset, count, cancellationToken).ConfigureAwait(false);
            long position = _position;
            _position = position + count;
            return new FileAssociatedImageItem(name, width, height, position, count, format, 0);
        }

        public async Task<FileAssociatedImageItem> WriteAssociatedImageItemAsync(Stream stream, string name, int width, int height, string format, CancellationToken cancellationToken = default)
        {
            if (stream is null)
            {
                throw new ArgumentNullException(nameof(stream));
            }

            if (format is null || format.Length != 4)
            {
                throw new ArgumentException("Format not supported.", nameof(format));
            }

            int bytesCoppied = 0;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(4096);
            try
            {
                while (true)
                {
                    int bytesRead = await stream.ReadAsync(buffer, 0, buffer.Length, cancellationToken).ConfigureAwait(false);
                    if (bytesRead == 0)
                    {
                        break;
                    }
                    await _stream.WriteAsync(buffer, 0, bytesRead, cancellationToken).ConfigureAwait(false);
                    bytesCoppied += bytesRead;
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

            long position = _position;
            _position = position + bytesCoppied;
            return new FileAssociatedImageItem(name, width, height, position, bytesCoppied, format, 0);
        }


        public Task<long> WriteAssociatedImagesBlockAsync(FileAssociatedImageItem[] items, CancellationToken cancellationToken = default)
            => WriteAssociatedImagesBlockAsync(string.Empty, items, 0, cancellationToken);

        public Task<long> WriteAssociatedImagesBlockAsync(FileAssociatedImageItem[] items, long nextBlockPointer, CancellationToken cancellationToken = default)
            => WriteAssociatedImagesBlockAsync(string.Empty, items, nextBlockPointer, cancellationToken);

        public async Task<long> WriteAssociatedImagesBlockAsync(string name, FileAssociatedImageItem[] items, long nextBlockPointer, CancellationToken cancellationToken = default)
        {
            if (name is null)
            {
                throw new ArgumentNullException(nameof(name));
            }
            if (items is null)
            {
                throw new ArgumentNullException(nameof(items));
            }

            int nameLength = Encoding.UTF8.GetByteCount(name);
            if (nameLength > 255)
            {
                throw new ArgumentException("Name too long.");
            }

            int totalLength = 3 + nameLength;
            foreach (FileAssociatedImageItem item in items)
            {
                totalLength += item.GetSize();
            }

            long position = _position;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(4 + totalLength + 8);
            try
            {
                BinaryPrimitives.WriteInt32LittleEndian(buffer.AsSpan(), totalLength);
                buffer[4] = (byte)nameLength;
                EncodingHelper.EncodeUtf8(name, buffer.AsSpan(5, nameLength));
                BinaryPrimitives.WriteInt16LittleEndian(buffer.AsSpan(5 + nameLength), (short)items.Length);
                await _stream.WriteAsync(buffer, 0, 7 + nameLength, cancellationToken).ConfigureAwait(false);
                _position += 7 + nameLength;

                int bytesWritten = 0;
                for (int i = 0; i < items.Length; i++)
                {
                    if (!items[i].TryWriteTo(buffer.AsSpan(bytesWritten), out int written))
                    {
                        throw new InvalidDataException();
                    }
                    bytesWritten += written;
                }
                BinaryPrimitives.WriteInt64LittleEndian(buffer.AsSpan(bytesWritten), nextBlockPointer);
                await _stream.WriteAsync(buffer, 0, bytesWritten + 8, cancellationToken).ConfigureAwait(false);
                _position += bytesWritten + 8;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

            return position;
        }

        public static int GetTileCollectionBlockLength(int tileCount)
        {
            return tileCount * FileContentRegion.MarshalSize;
        }

        public async Task<long> SkipTileCollectionAsync(int tileCount, CancellationToken cancellationToken = default)
        {
            long position = _position;
            const int batchSize = 4096 / FileContentRegion.MarshalSize;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(4096);
            try
            {
                Array.Clear(buffer, 0, 4096);

                int filled = 0;
                for (int i = 0; i < tileCount; i++)
                {
                    filled++;
                    if (filled == batchSize)
                    {
                        await _stream.WriteAsync(buffer, 0, filled * FileContentRegion.MarshalSize, cancellationToken).ConfigureAwait(false);
                        _position += filled * FileContentRegion.MarshalSize;
                        filled = 0;
                    }
                }

                if (filled > 0)
                {
                    await _stream.WriteAsync(buffer, 0, filled * FileContentRegion.MarshalSize, cancellationToken).ConfigureAwait(false);
                    _position += filled * FileContentRegion.MarshalSize;
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
            return position;
        }

        public async Task<long> WriteTileCollectionAsync(IEnumerable<FileContentRegion> tiles, CancellationToken cancellationToken = default)
        {
            if (tiles is null)
            {
                throw new ArgumentNullException(nameof(tiles));
            }

            long position = _position;
            const int batchSize = 4096 / FileContentRegion.MarshalSize;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(4096);
            try
            {
                int filled = 0;
                foreach (FileContentRegion tile in tiles)
                {
                    tile.TryWriteTo(buffer.AsSpan(filled * FileContentRegion.MarshalSize, FileContentRegion.MarshalSize));
                    filled++;

                    if (filled == batchSize)
                    {
                        await _stream.WriteAsync(buffer, 0, filled * FileContentRegion.MarshalSize, cancellationToken).ConfigureAwait(false);
                        _position += filled * FileContentRegion.MarshalSize;
                        filled = 0;
                    }
                }

                if (filled > 0)
                {
                    await _stream.WriteAsync(buffer, 0, filled * FileContentRegion.MarshalSize, cancellationToken).ConfigureAwait(false);
                    _position += filled * FileContentRegion.MarshalSize;
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
            return position;
        }

        public static int GetLevelCollectionBlockLength(int levelCount)
        {
            return levelCount * 8;
        }

        public async Task<long> SkipLevelCollectionAsync(int levelCount, CancellationToken cancellationToken = default)
        {
            long position = _position;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(levelCount * 8);
            try
            {
                Array.Clear(buffer, 0, levelCount * 8);
                await _stream.WriteAsync(buffer, 0, levelCount * 8, cancellationToken).ConfigureAwait(false);
                _position += levelCount * 8;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
            return position;
        }

        public async Task<long> WriteLevelCollectionAsync(long[] tileCollectionOffset, CancellationToken cancellationToken = default)
        {
            if (tileCollectionOffset is null)
            {
                throw new ArgumentNullException(nameof(tileCollectionOffset));
            }

            long position = _position;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(tileCollectionOffset.Length * 8);
            try
            {
                for (int i = 0; i < tileCollectionOffset.Length; i++)
                {
                    long offset = tileCollectionOffset[i];
                    BinaryPrimitives.WriteInt64LittleEndian(buffer.AsSpan(i * 8, 8), offset);
                }
                await _stream.WriteAsync(buffer, 0, tileCollectionOffset.Length * 8, cancellationToken).ConfigureAwait(false);
                _position += tileCollectionOffset.Length * 8;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
            return position;
        }

        public static int GetMainBlockLength()
        {
            return FileMainBlock.MarshalSize;
        }

        public async Task<long> SkipMainBlockAsync(CancellationToken cancellationToken = default)
        {
            long position = _position;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(FileMainBlock.MarshalSize);
            try
            {
                Array.Clear(buffer, 0, FileMainBlock.MarshalSize);
                await _stream.WriteAsync(buffer, 0, FileMainBlock.MarshalSize, cancellationToken).ConfigureAwait(false);
                _position = position + FileMainBlock.MarshalSize;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
            return position;
        }

        public async Task<long> WriteMainBlockAsync(FileMainBlock mainBlock, CancellationToken cancellationToken = default)
        {
            long position = _position;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(FileMainBlock.MarshalSize);
            try
            {
                mainBlock.TryWriteTo(buffer);
                await _stream.WriteAsync(buffer, 0, FileMainBlock.MarshalSize, cancellationToken).ConfigureAwait(false);
                _position = position + FileMainBlock.MarshalSize;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
            return position;
        }
    }
}
